package com.company.并查集;

/**
 * 并查集，手写并查集（值得记忆，万能）
 * 受限条件下可到达节点的数目
 * 关键：忽略受限的点;细节上尽量不要乱更改,isRestricted的Boolean不必改为int
 */
public class LeetCode2368 {
    private int[] isRestricted;
    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        isRestricted=new int[n];
        for (int i = 0; i < restricted.length; i++) {
            isRestricted[restricted[i]]=1;
        }
        UnionFind unionFind=new UnionFind(n);
        for (int i = 0; i < edges.length; i++) {
            if(isOK(edges[i][0],edges[i][1])){
                unionFind.merge(edges[i][0],edges[i][1]);
            }
        }
        return unionFind.count(0);
    }

    private Boolean isOK(int x,int y){
        if (isRestricted[x]!=1&&isRestricted[y]!=1)
            return true;
        else
            return false;
    }

    class UnionFind {
        private int[] f;
        private int[] rank;//深度

        public UnionFind(int n) {
            f = new int[n];
            rank = new int[n];
            for (int i = 0; i < n; i++) {
                f[i] = i;
            }
        }

        public void merge(int x, int y) {
            int rx = find(x);
            int ry = find(y);
            if (rx != ry) {
                if (rank[rx] > rank[ry]) {
                    f[ry] = rx;
                } else if (rank[rx] < rank[ry]) {
                    f[rx] = ry;
                } else {
                    f[ry] = rx;
                    rank[rx]++;
                }
            }
        }

        public int find(int x) {
            if (x != f[x]) {
                x = find(f[x]);
            }
            return f[x];
        }

        public int count(int x) {
            int cnt = 0;
            int rt = find(x);
            for (int i = 0; i < f.length; i++) {
                if (rt == find(i)) {
                    cnt++;
                }
            }
            return cnt;
        }
    }
}
